# -*- coding: utf-8 -*-
# @Author  : gaoyu
# @Time    : 2024/9/26
# @Function: 一对一密钥分享方案实现

import base64
import json
import os
import random
import sys
from typing import List, Dict, Any

from gmssl import sm4
from loguru import logger

# 配置loguru日志格式，使其更加花里胡哨
logger.remove()
logger.add(
    sys.stderr,
    format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
           "<level>{level: <8}</level> | "
           "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> | "
           "<level>{message}</level>",
    level="INFO",
    colorize=True,
)


def log_section(title):
    logger.info(f"\n{'=' * 20} 🔐 {title} 🔐 {'=' * 20}")


def log_subheader(title):
    logger.info(f"\n{'-' * 15} 🔹 {title} 🔹 {'-' * 15}")


def log_header(title):
    logger.info(f"\n{'=' * 50}")
    logger.info(f"{'=' * 15} 🌟 {title} 🌟 {'=' * 15}")
    logger.info(f"{'=' * 50}")


class PairwiseCrypto:
    """
    阈值密码学方案实现一对一密钥分享方案实现
    允许数据拥有者与每位参与者单独分享密钥
    """

    def __init__(self):
        """初始化一对一密钥分享工具"""
        self.sm4_key = None  # 用于加密实际数据的对称密钥
        self.encrypted_data = None  # 存储加密后的数据
        self.encrypted_data_base64 = None  # Base64编码的加密数据
        self.owner_shares = {}  # 存储数据拥有者持有的密钥分片
        self.participant_shares = {}  # 存储各参与者持有的密钥分片

    def generate_sm4_key(self) -> bytes:
        """生成随机SM4密钥"""
        return bytes([random.randint(0, 255) for _ in range(16)])

    def split_key_for_pair(self, key: bytes) -> tuple:
        """
        将密钥分成两份，需要两份都有才能重构
        
        Args:
            key: 要分割的密钥
            
        Returns:
            tuple: (share1, share2) 两个密钥分片
        """
        # 生成随机掩码
        mask = bytes([random.randint(0, 255) for _ in range(len(key))])

        # 第一份是随机掩码
        share1 = mask

        # 第二份是原始密钥与掩码的异或结果
        share2 = bytes([key[i] ^ mask[i] for i in range(len(key))])

        return share1, share2

    def reconstruct_key(self, share1: bytes, share2: bytes) -> bytes:
        """
        从两个密钥分片重构密钥
        
        Args:
            share1: 第一个密钥分片
            share2: 第二个密钥分片
            
        Returns:
            bytes: 重构的密钥
        """
        # 通过异或操作重构密钥
        return bytes([share1[i] ^ share2[i] for i in range(len(share1))])

    def set_sm4_key(self, key: bytes) -> None:
        """设置SM4对称密钥"""
        self.sm4_key = key if key is not None else self.generate_sm4_key()

    def exist_encrypted_package(self, key_file: str) -> bool:
        """获取已加密的数据包"""
        # 尝试从文件加载已加密的数据
        if key_file and os.path.exists(key_file):
            try:
                with open(key_file, 'r') as f:
                    key_data = json.load(f)
                    if 'encrypted_data' in key_data:
                        self.encrypted_data = base64.b64decode(key_data['encrypted_data'])
                        self.encrypted_data_base64 = key_data['encrypted_data']
                        logger.info(f"✅ 从文件 {key_file} 成功加载已加密数据")
                        return True
            except Exception as e:
                logger.warning(f"⚠️ 从文件加载加密数据失败: {e}，将重新加密数据")
                return False

    def encrypt_data(self, data: str, key_file: str = None) -> None:
        """
        仅加密数据，不分配密钥分片
        
        Args:
            data: 要加密的数据
            key_file: 可选的密钥文件路径，如果提供则尝试从文件读取/保存加密数据
        """
        # 尝试从文件加载已加密的数据
        if key_file and os.path.exists(key_file):
            try:
                with open(key_file, 'r') as f:
                    key_data = json.load(f)
                    if 'encrypted_data' in key_data:
                        self.encrypted_data = base64.b64decode(key_data['encrypted_data'])
                        self.encrypted_data_base64 = key_data['encrypted_data']
                        logger.info(f"✅ 从文件 {key_file} 成功加载已加密数据")
                        return
            except Exception as e:
                logger.warning(f"⚠️ 从文件加载加密数据失败: {e}，将重新加密数据")

        # 使用提供的密钥或生成新的SM4对称密钥
        self.sm4_key = self.sm4_key if self.sm4_key is not None else self.generate_sm4_key()

        # 使用SM4加密数据
        sm4_crypt = sm4.CryptSM4()
        sm4_crypt.set_key(self.sm4_key, sm4.SM4_ENCRYPT)

        # 确保数据长度是16的倍数（SM4块大小），添加PKCS#7填充
        data_bytes = data.encode('utf-8')
        padding_length = 16 - (len(data_bytes) % 16)
        if padding_length < 16:
            data_bytes += bytes([padding_length]) * padding_length

        # 加密数据
        self.encrypted_data = sm4_crypt.crypt_ecb(data_bytes)
        self.encrypted_data_base64 = base64.b64encode(self.encrypted_data).decode('utf-8')

        # 如果提供了文件路径，保存加密数据
        if key_file:
            try:
                key_data = {}
                if os.path.exists(key_file):
                    with open(key_file, 'r') as f:
                        key_data = json.load(f)

                key_data['encrypted_data'] = self.encrypted_data_base64

                with open(key_file, 'w') as f:
                    json.dump(key_data, f, indent=2)
                logger.info(f"✅ 加密数据已保存到文件 {key_file}")
            except Exception as e:
                logger.warning(f"⚠️ 保存加密数据到文件失败: {e}")

    def set_encrypted_data(self, encrypted_data: bytes) -> None:
        """
        设置已加密的数据
        
        Args:
            encrypted_data: 已加密的数据字节
        """
        self.encrypted_data = encrypted_data
        self.encrypted_data_base64 = base64.b64encode(encrypted_data).decode('utf-8')

    def get_encrypted_data(self) -> bytes:
        """
        获取已加密的数据字节
        
        Returns:
            bytes: 加密后的数据字节
        """
        if not self.encrypted_data:
            raise ValueError("请先调用encrypt_data方法加密数据")
        return self.encrypted_data

        # 清空之前的密钥分片
        # self.owner_shares = {}
        # self.participant_shares = {}

    def add_participant(self, participant_id: str) -> Dict[str, str]:
        """
        为单个参与者生成密钥分片
        
        Args:
            participant_id: 参与者ID
            
        Returns:
            Dict: 包含该参与者的密钥分片信息
        """
        if not self.sm4_key:
            raise ValueError("请先调用encrypt_data方法加密数据")

        # 为参与者单独分割密钥
        owner_share, participant_share = self.split_key_for_pair(self.sm4_key)

        # 存储分片
        owner_share_b64 = base64.b64encode(owner_share).decode('utf-8')
        participant_share_b64 = base64.b64encode(participant_share).decode('utf-8')

        self.owner_shares[participant_id] = owner_share_b64
        self.participant_shares[participant_id] = participant_share_b64

        # 返回该参与者的分片信息
        return {
            'participant_id': participant_id,
            'owner_share': owner_share_b64,
            'participant_share': participant_share_b64
        }

    def get_encrypted_package(self) -> Dict[str, Any]:
        """
        获取当前的加密包
        
        Returns:
            Dict: 包含加密数据和所有参与者的密钥分片
        """
        if not self.encrypted_data:
            raise ValueError("请先调用encrypt_data方法加密数据")

        return {
            'encrypted_data': self.encrypted_data_base64,
            'owner_shares': self.owner_shares,
            'participant_shares': self.participant_shares
        }

    def encrypt_data_for_multiple_pairs(self, data: str, participants: List[str], key: bytes = None) -> Dict[str, Any]:
        """
        加密数据并为多个参与者分割密钥（向后兼容）
        
        Args:
            data: 要加密的数据
            participants: 参与者列表
            key: 可选的预设SM4密钥，如果不提供则自动生成
            
        Returns:
            Dict: 包含加密数据和每对参与者的密钥分片
        """
        # 先加密数据
        self.encrypt_data(data, key)

        # 为每个参与者生成密钥分片
        for participant in participants:
            self.add_participant(participant)

        # 返回完整的加密包
        return self.get_encrypted_package()

    def decrypt_data(self, encrypted_package: Dict[str, Any], owner_share: str, participant_share: str) -> str:
        """
        使用拥有者和参与者的密钥分片解密数据
        
        Args:
            encrypted_package: 包含加密数据的字典
            owner_share: 拥有者的密钥分片
            participant_share: 参与者的密钥分片
            
        Returns:
            str: 解密后的数据
        """
        try:
            # 解码密钥分片
            owner_share_bytes = base64.b64decode(owner_share)
            participant_share_bytes = base64.b64decode(participant_share)

            # 重构密钥
            reconstructed_key = self.reconstruct_key(owner_share_bytes, participant_share_bytes)

            # 使用重构的密钥解密数据
            encrypted_data = base64.b64decode(encrypted_package['encrypted_data'])
            sm4_crypt = sm4.CryptSM4()
            sm4_crypt.set_key(reconstructed_key, sm4.SM4_DECRYPT)
            decrypted_data = sm4_crypt.crypt_ecb(encrypted_data)

            # 移除PKCS#7填充
            if decrypted_data:
                padding_length = decrypted_data[-1]
                if padding_length < 16 and all(b == padding_length for b in decrypted_data[-padding_length:]):
                    decrypted_data = decrypted_data[:-padding_length]

            # 解码为字符串
            result = decrypted_data.decode('utf-8')

            # 检查解密结果是否为空
            if not result:
                raise ValueError("解密结果为空字符串")

            return result
        except Exception as e:
            logger.error(f"❌ 解密过程中出错: {e}")
            raise
